#!/bin/sh

# Copyright 2024 Balena Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#
# Update the TPM2 sealing policy on commit.
#
# This updates the passphrase to use the newly computed value of PCR7 if any
# measured variables have been changed, such as by appending to dbx during HUP.
#
# Additonally, the hook will migrate from systems that were setup to seal the
# LUKS passphrase using PCRs 0,1,2,3 to using 0,2,3,7.
#
# This is done in a commit hook to preserve fallback capability until the new
# OS is healthy.
#

set -o errexit

# shellcheck disable=SC1091
. /usr/libexec/os-helpers-sb

EFI_MOUNT_DIR="/mnt/efi"
PASSPHRASE_FILE="$(mktemp -t)"
SESSION_CTX="$(mktemp -t)"
CURRENT_POLICY_PATH="$(find "${EFI_MOUNT_DIR}" -type d -name "policies.*")"
tpm2_startauthsession --policy-session -S "${SESSION_CTX}"
tpm2_policypcr -S "${SESSION_CTX}" -l "sha256:0,2,3,7"

update_reason=""
POLICIES="$(find "${CURRENT_POLICY_PATH}" -type f | sort | xargs)"
if [ "$(echo "${POLICIES}" | wc -w)" -gt 1 ]; then
	tpm2_policyor -S "${SESSION_CTX}" "sha256:$(echo "${POLICIES}" | sed 's/ /,/g')"
	update_reason="Combined policy in use"
fi

trap 'tpm2_flushcontext "${SESSION_CTX}" 2>/dev/null' EXIT

if tpm_nvram_retrieve_passphrase "session:${SESSION_CTX}" "${PASSPHRASE_FILE}"; then
	echo "Retrieved passphrase from TPM NVRAM"
	rm -f "${EFI_MOUNT_DIR}/balena-luks.*"
else
	echo "Failed to unlock passphrase, abort"
	exit 1
fi

tpm2_flushcontext "${SESSION_CTX}" >/dev/null 2>&1

POLICY_PATH="$(mktemp -t -d policies.XXXXX)"
POLICY="$(mktemp -p "${POLICY_PATH}")"
PCRS="0,2,3,7"
PCR_VAL_BIN="$(mktemp -t)"
EFI_BINARIES=" \
	$(find "${EFI_MOUNT_DIR}" -name bootx64.efi -print -quit) \
	$(find /boot -name bzImage -print -quit)
	"

if [ "$(firmware_measures_efibins)" = "measured" ]; then
	generate_pcr_digests \
		"${PCRS}" \
		"${PCR_VAL_BIN}" \
		"use_pcr2" \
		"${EFI_BINARIES}"
else
	generate_pcr_digests \
		"${PCRS}" \
		"${PCR_VAL_BIN}" \
		"use_pcr2"
fi

stdout=/proc/self/fd/1
current_digest="$(tpm2_pcrread --quiet "sha256:7" -o "${stdout}" | _hexencode)"
computed_digest="$(dd if="${PCR_VAL_BIN}" bs=1 status=none skip=$((32 * 3)) | _hexencode)"
if [ "${current_digest}" != "${computed_digest}" ]; then
	update_reason="PCR7 computed value changed"
fi

if [ -n "${update_reason}" ]; then
	echo "${update_reason}, updating policy"
	print_pcr_val_bin "${PCRS}" "${PCR_VAL_BIN}"

	tpm2_createpolicy --policy-pcr \
		-l "sha256:${PCRS}" \
		-f "${PCR_VAL_BIN}" \
		-L "${POLICY}"
	
	tpm_nvram_store_passphrase "${PASSPHRASE_FILE}" "${POLICY_PATH}"

	# shellcheck disable=SC2012
	if [ "$(ls -1 "${POLICY_PATH}" | wc -l)" -gt 1 ]; then
		cp -rf "${POLICY_PATH}" "${EFI_MOUNT_DIR}"
	fi

	rm -rf "${PCR_VAL_BIN}" \
	      "${PASSPHRASE_FILE}" \
	      "${CURRENT_POLICY_PATH}"

	sync

	if [ "${current_digest}" != "${computed_digest}" ]; then
		# reboot to ensure the passphrase can be unlocked again,
		# otherwise HUP won't work until the device is manually
		# rebooted
		reboot
	fi
else
	echo "PCR7 computed value is unchanged"
fi
